﻿using System;
using System.Collections.Generic;
using System.Text;
using System.IO;
using MPI;

namespace Dwarf.Managed
{
    /// <summary>
    /// Solve the Sparse Linear Algebra problem with MPI.NET parallelization.
    /// </summary>
    /// <remarks> 
    /// For the demonstration purposes of the project, 
    /// we have chosen sparse matrix-vector multiplication.
    /// </remarks>
    class Solver
    {
        /// <summary>
        /// Number of rows in sparse matrix
        /// </summary>
        public int rowNum;
        /// <summary>
        /// Number of non-zero values in sparse matrix
        /// </summary>
        public int valNum;

        /// <summary>
        /// sparse matrix in Compressed Row Storage (CRS) format
        /// </summary>
        public double[] val;
        public int[] col_ind;
        public int[] row_ptr;

        /// <summary>
        /// Vector
        /// </summary>
        public double[] vector;
        /// <summary>
        /// Result vector
        /// </summary>
        public double[] result;

        /// <summary>
        /// Empty constructor.
        /// </summary>
        public Solver()
        {
        }

        /// <summary>
        /// MPI.NET based method for sparse matrix-vector multiplication.
        /// </summary>
        public void Solve()
        {
            bool isRootProcess = Communicator.world.Rank == Settings.ROOT_PROCESS;
            int numProcesses = Math.Min(Communicator.world.Size, rowNum);

            double[] valLocal = null;
            int[] col_indLocal = null;
            int[] row_ptrLocal = null;

            int[] numRowsOnProc = new int[Communicator.world.Size];                        
            int rowCount = rowNum;
            for (int i = 0; i < numProcesses; i++)
            {
                numRowsOnProc[i] = Math.Min(rowNum / numRowsOnProc.Length, rowCount);
                rowCount -= numRowsOnProc[i];
            }
            int numRowsLocal = numRowsOnProc[Communicator.world.Rank];
            
            Communicator.world.ScatterFromFlattened(row_ptr, numRowsOnProc, Settings.ROOT_PROCESS, ref row_ptrLocal);

            int[] lastRow_ptr = new int[numProcesses]; 
            int[] lastRow_ptrLocal = null;
            if (isRootProcess)
            {
                int lastPtr = 0;
                for (int i = 0; i < numProcesses; i++)
                {
                    lastRow_ptr[i] = row_ptr[lastPtr + numRowsOnProc[i]];
                    lastPtr += numRowsOnProc[i];
                }
            }

            Communicator.world.ScatterFromFlattened(lastRow_ptr, 1, Settings.ROOT_PROCESS, ref lastRow_ptrLocal);

            int[] numNonZeroElemsOnProc = new int[Communicator.world.Size];
            if (isRootProcess)
            {
                int rowPtr = 0;
                for (int i = 0; i < numNonZeroElemsOnProc.Length; i++)
                {
                    numNonZeroElemsOnProc[i] = row_ptr[rowPtr + numRowsOnProc[i]] - row_ptr[rowPtr];
                    rowPtr += numRowsOnProc[i];
                }
            }

            Communicator.world.Broadcast(ref numNonZeroElemsOnProc, Settings.ROOT_PROCESS);

            Communicator.world.ScatterFromFlattened(col_ind, numNonZeroElemsOnProc, Settings.ROOT_PROCESS, ref col_indLocal);
            Communicator.world.ScatterFromFlattened(val, numNonZeroElemsOnProc, Settings.ROOT_PROCESS, ref valLocal);

            double[] tempResult = new double[numRowsLocal];
            for (int i = 0; i < numRowsLocal; i++)                                //Loop for number of rows in matrix
            {
                //The C# compler/runtime is not
                //hoisting the array addressing out of the loop
                //hence the following makes a big difference
                //result[i] = 0;
                var cell = 0.0;
                var ptr = (i == numRowsLocal - 1 ? lastRow_ptrLocal[0] : row_ptrLocal[i + 1]) - row_ptrLocal[0];
                for (int j = row_ptrLocal[i] - row_ptrLocal[0]; j < ptr; j++)  //Loop for number of non-zero elements in current row
                {
                    cell += vector[col_indLocal[j] - 1] * valLocal[j];
                }
                //see above comment
                tempResult[i] = cell;
            }

            result = Communicator.world.GatherFlattened(tempResult, Settings.ROOT_PROCESS);
        }
    }
}
